{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "ej-jaFpE_2_B"
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "import pickle\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "MAX_VOCAB_SIZE = 10000  # 词表长度限制\n",
    "UNK, PAD = '<UNK>', '<PAD>'  # 未知字 padding符号\n",
    "MIN_FREQ = 1  # 出现频率大于该参数的字才放入词典"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "gXQz8b4IASnS"
   },
   "outputs": [],
   "source": [
    "def build_vocab(train_path, class_ls_path):  # 构建词典（默认是字符级）\n",
    "    vocab_dic = {}  # 字典\n",
    "    class_set = set()  # 集合\n",
    "    with open(train_path, 'r', encoding='utf-8') as f:\n",
    "        for line in tqdm(f):\n",
    "            lin = line.strip()  # 去除头尾的空格 换行符 制表符\n",
    "            if not lin:\n",
    "                continue\n",
    "            content, label = lin.split()\n",
    "            vocab_dic[content] = vocab_dic.get(content, 0) + 1  # 每个字计数\n",
    "            class_set.add(label)\n",
    "        vocab_ls = sorted([_ for _ in vocab_dic.items() if _[1] >= MIN_FREQ],\n",
    "                          key=lambda x: x[1], reverse=True)[:MAX_VOCAB_SIZE]\n",
    "                          \n",
    "        class_ls = list(sorted(class_set))\n",
    "        with open(class_ls_path, 'w', encoding='utf-8') as cf:\n",
    "            cf.write('\\n'.join(str(label) for label in class_ls))\n",
    "            cf.write('\\n' + PAD )\n",
    "\n",
    "        vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_ls)}\n",
    "        vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})  # 将UNK和PAD索引加到词典最后\n",
    "    return vocab_dic, vocab_dic[PAD], len(class_ls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "75-zNEYQDBXm",
    "outputId": "e9a12b97-a8c9-4b68-b106-b3371b05c7e0"
   },
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'ner.train'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mFileNotFoundError\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[3], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m vocab_dic, word_pad_id, label_pad_id \u001B[38;5;241m=\u001B[39m \u001B[43mbuild_vocab\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43mner.train\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43mner.label\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m)\u001B[49m\n",
      "Cell \u001B[1;32mIn[2], line 4\u001B[0m, in \u001B[0;36mbuild_vocab\u001B[1;34m(train_path, class_ls_path)\u001B[0m\n\u001B[0;32m      2\u001B[0m vocab_dic \u001B[38;5;241m=\u001B[39m {}  \u001B[38;5;66;03m# 字典\u001B[39;00m\n\u001B[0;32m      3\u001B[0m class_set \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mset\u001B[39m()  \u001B[38;5;66;03m# 集合\u001B[39;00m\n\u001B[1;32m----> 4\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m \u001B[38;5;28;43mopen\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mtrain_path\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43mr\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mencoding\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43mutf-8\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m)\u001B[49m \u001B[38;5;28;01mas\u001B[39;00m f:\n\u001B[0;32m      5\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m line \u001B[38;5;129;01min\u001B[39;00m tqdm(f):\n\u001B[0;32m      6\u001B[0m         lin \u001B[38;5;241m=\u001B[39m line\u001B[38;5;241m.\u001B[39mstrip()  \u001B[38;5;66;03m# 去除头尾的空格 换行符 制表符\u001B[39;00m\n",
      "File \u001B[1;32mD:\\ananaconda\\envs\\pytorch\\lib\\site-packages\\IPython\\core\\interactiveshell.py:284\u001B[0m, in \u001B[0;36m_modified_open\u001B[1;34m(file, *args, **kwargs)\u001B[0m\n\u001B[0;32m    277\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m file \u001B[38;5;129;01min\u001B[39;00m {\u001B[38;5;241m0\u001B[39m, \u001B[38;5;241m1\u001B[39m, \u001B[38;5;241m2\u001B[39m}:\n\u001B[0;32m    278\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\n\u001B[0;32m    279\u001B[0m         \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mIPython won\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mt let you open fd=\u001B[39m\u001B[38;5;132;01m{\u001B[39;00mfile\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m by default \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    280\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mas it is likely to crash IPython. If you know what you are doing, \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    281\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124myou can use builtins\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m open.\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    282\u001B[0m     )\n\u001B[1;32m--> 284\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m io_open(file, \u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n",
      "\u001B[1;31mFileNotFoundError\u001B[0m: [Errno 2] No such file or directory: 'ner.train'"
     ]
    }
   ],
   "source": [
    "vocab_dic, word_pad_id, label_pad_id = build_vocab('ner.train', 'ner.label')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Wj57w63dDJng",
    "outputId": "0d359844-cc76-4baf-ba8b-d8609b3446bb"
   },
   "outputs": [],
   "source": [
    "vocab_dic['专']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "72ri1TsWRlmu",
    "outputId": "3278f20c-eb26-4c97-cf7f-d0fcffa95705"
   },
   "outputs": [],
   "source": [
    "word_pad_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PTZGYBxvRm5k",
    "outputId": "9d6cca7a-6359-4a57-abd5-eba0e8eae955"
   },
   "outputs": [],
   "source": [
    "label_pad_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FKov1SHj8mld"
   },
   "outputs": [],
   "source": [
    "def build_corpus(filepath):\n",
    "    word_lists = []\n",
    "    tag_lists = []\n",
    "    with open(filepath, 'r', encoding='utf-8') as f:\n",
    "        word_list = []\n",
    "        tag_list = []\n",
    "        for line in f:\n",
    "          lin = line.strip()   \n",
    "          if  lin:\n",
    "              word, tag = line.strip().split()\n",
    "              word_list.append(word)\n",
    "              tag_list.append(tag)\n",
    "          else:\n",
    "              word_lists.append(word_list)\n",
    "              tag_lists.append(tag_list)\n",
    "              word_list = []\n",
    "              tag_list = []\n",
    "\n",
    "    word_lists = sorted(word_lists, key=lambda x: len(x), reverse=True)\n",
    "    tag_lists = sorted(tag_lists, key=lambda x: len(x), reverse=True)\n",
    "    return word_lists, tag_lists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FGMaRzyG9tCR"
   },
   "outputs": [],
   "source": [
    "word_lists, tag_lists = build_corpus('ner.train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gqUuwybU_Ax6"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 514
    },
    "id": "SosgzCy2_hkv",
    "outputId": "eeb3eed1-2cb5-43bc-c2d2-6244e435021b"
   },
   "outputs": [],
   "source": [
    "sen_lengths = [len(sent) for sent in word_lists]\n",
    "plt.figure(figsize=(15, 5))\n",
    "plt.hist(sen_lengths) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BMcRNNIz_Hpx"
   },
   "outputs": [],
   "source": [
    "class Config(object):\n",
    "    def __init__(self):\n",
    "        self.train_path = 'ner.train'  # *\n",
    "        self.dev_path = 'ner.dev'  # *\n",
    "        self.class_ls_path = 'ner.label'  # *\n",
    "        self.pretrain_dir = './'  # 前期下载的预训练词向量*\n",
    "        self.test_path = 'ner.test'   \n",
    "        self.vocab_path = 'vocab.pkl'\n",
    "        self.model_save_dir = 'checkpoint'\n",
    "        self.model_save_name = self.model_save_dir + '/BiLSTM_CRF.ckpt'  # 保存最佳dev acc模型\n",
    "        self.max_seq_len = 200\n",
    "        # 可调整的参数\n",
    "        # 搜狗新闻:embedding_SougouNews.npz, 腾讯:embedding_Tencent.npz,  若不存在则后期生成\n",
    "        # 随机初始化:random\n",
    "        self.embedding_type = 'embedding_SougouNews.npz'\n",
    "        self.use_gpu = True  # 是否使用gpu(有则加载 否则自动使用cpu)\n",
    "        self.batch_size = 64\n",
    "        self.num_epochs = 100  # 训练轮数\n",
    "        self.num_workers = 8  # 启用多线程\n",
    "        self.learning_rate = 0.001  # 训练发现0.001比0.01收敛快(Adam)\n",
    "        self.embedding_dim = 300  # 词嵌入维度\n",
    "        self.hidden_size = 300  # 隐藏层维度\n",
    "        self.num_layers = 2  # RNN层数\n",
    "        self.bidirectional = True  # 双向 or 单向\n",
    "        self.require_improvement = 10  # 1个epoch若在dev上acc未提升则自动结束\n",
    "\n",
    "        # 由前方参数决定  不用修改\n",
    "        self.class_ls = []\n",
    "        self.num_class = len(self.class_ls)\n",
    "        self.vocab_len = 0  # 词表大小(训练集总的字数(字符级)） 在embedding层作为参数 后期赋值\n",
    "        self.embedding_pretrained = None  # 根据config.embedding_type后期赋值  random:None  else:tensor from embedding_type\n",
    "        if self.use_gpu and torch.cuda.is_available():\n",
    "            self.device = 'cuda:0'\n",
    "        else:\n",
    "            self.device = 'cpu'\n",
    "\n",
    "class NERDataset(Dataset):\n",
    "    def __init__(self, sentences, tags, vocab, config):\n",
    "        self.sentences = sentences\n",
    "        self.tags = tags\n",
    "        self.label_dic = self._getLabelDic(config)\n",
    "        self.vocab = vocab\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.sentences)\n",
    "    \n",
    "    def __getitem__(self, item):\n",
    "        sentence = self.sentences[item]\n",
    "        tag = self.tags[item]\n",
    "        seq_len = len(sentence)\n",
    "\n",
    "        # convert the sentences and tags into numerical format\n",
    "        word_tokens = [self.vocab.get(word, self.vocab.get(UNK)) for word in sentence]\n",
    "        tag_tokens = [self.label_dic[t] for t in tag]\n",
    "\n",
    "        \n",
    "        return torch.LongTensor(word_tokens),  torch.LongTensor(tag_tokens)\n",
    "\n",
    "    def _getLabelDic(self, config):\n",
    "        label_dic ={}\n",
    "        with open(config.class_ls_path, 'r', encoding='utf-8') as f:\n",
    "            for idx, line in enumerate(f):\n",
    "                label = line.strip()\n",
    "                label_dic[label] = idx\n",
    "        return label_dic\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ggv6_DXJDBFg"
   },
   "outputs": [],
   "source": [
    "config = Config()  # 实例化各参数\n",
    "\n",
    "train_dataset = NERDataset(word_lists, tag_lists, vocab_dic, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uSp13RpREIW3",
    "outputId": "897ea0dd-5857-48b1-c460-699c01ed0b7f"
   },
   "outputs": [],
   "source": [
    "train_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RZmk0I8hELYm"
   },
   "outputs": [],
   "source": [
    "def collate_fn(data):\n",
    "    data.sort(key=lambda x: len(x[0]), reverse=True)\n",
    "    sentences, tags = zip(*data)\n",
    "\n",
    "    # Merge questions (from tuple of 1D tensor to 2D tensor).\n",
    "    sent_lengths = [len(sent) for sent in sentences]\n",
    "    inputs = torch.full( (len(sentences), max(sent_lengths)), word_pad_id ).long() # vocab_dic['PAD] torch.full\n",
    "    labels = torch.full( (len(sentences), max(sent_lengths)), label_pad_id  ).long() # label_dic['O']. torch.full\n",
    "    \n",
    "    for i, (sent, lab) in enumerate(zip(sentences, tags)):\n",
    "        end = sent_lengths[i]\n",
    "        inputs[i, :end] = sent[:end]\n",
    "        labels[i, :end] = lab[:end]\n",
    "    return inputs, labels, sent_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w796Af2tEUuS"
   },
   "outputs": [],
   "source": [
    "train_data_loader =  DataLoader(train_dataset, 128, shuffle=True, collate_fn=collate_fn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fJzrqre-EyAS",
    "outputId": "606954c7-75cb-4a1b-b702-96d70bf35f61"
   },
   "outputs": [],
   "source": [
    "next(iter(train_data_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8VxN6qlgDqAe"
   },
   "outputs": [],
   "source": [
    "def build_dataset(config):\n",
    "    \"\"\"\n",
    "    函数说明： 根据训练、验证、测试数据返回其由索引组成的张量\n",
    "\n",
    "    Parameter：\n",
    "    ----------\n",
    "        config - ...\n",
    "    Return:\n",
    "    -------\n",
    "        train_loader - 用于模型输入的训练集DataLoader\n",
    "        dev_loader - 用于模型输入的验证集DataLoader\n",
    "        test_loader - 用于模型输入的测试集DataLoader\n",
    "\n",
    "    \"\"\"\n",
    "    if os.path.exists(config.vocab_path):  # 加载词典\n",
    "        vocab = pickle.load(open(config.vocab_path, 'rb'))\n",
    "    else:\n",
    "        vocab, word_pad_id, label_pad_id = build_vocab(config.train_path, config.class_ls_path)  # 用训练数据构建词典\n",
    "        with open(config.vocab_path,'wb') as f:\n",
    "            pickle.dump(vocab, f)  # 存储每个字及对应索引的字典 eg：我：56  vocab['我’]=56\n",
    "    config.vocab_len = len(vocab)\n",
    "    config.class_ls = [x.strip() for x in open(config.class_ls_path, 'r', encoding='utf-8').readlines()]\n",
    "    print('\\nVocab size: {}'.format(len(vocab)))\n",
    "\n",
    "\n",
    "    word_lists, tag_lists = build_corpus(config.train_path)\n",
    "    train_data = NERDataset(word_lists, tag_lists, vocab_dic, config)\n",
    "\n",
    "    word_lists, tag_lists = build_corpus(config.dev_path)\n",
    "    dev_data = NERDataset(word_lists, tag_lists, vocab_dic, config)\n",
    "    train_loader = DataLoader(dataset=train_data,\n",
    "                              batch_size=config.batch_size,\n",
    "                              shuffle=True,\n",
    "                              num_workers=config.num_workers,\n",
    "                              collate_fn=collate_fn)\n",
    "    dev_loader = DataLoader(dataset=dev_data,\n",
    "                            batch_size=config.batch_size,\n",
    "                            shuffle=False,\n",
    "                            num_workers=config.num_workers,\n",
    "                            collate_fn=collate_fn)\n",
    "\n",
    "    \n",
    "    if os.path.exists(config.test_path):\n",
    "        word_lists, tag_lists = build_corpus(config.test_path)\n",
    "        test_data = NERDataset(word_lists, tag_lists, vocab_dic, config)\n",
    "        test_loader = DataLoader(dataset=test_data,\n",
    "                                  batch_size=config.batch_size,\n",
    "                                  shuffle=False,\n",
    "                                  num_workers=config.num_workers,\n",
    "                                 collate_fn=collate_fn)\n",
    "    else:  # 若无测试数据则加载验证集进行最终测试\n",
    "        test_loader = dev_loader\n",
    "    config.embedding_pretrained = torch.tensor(extract_vocab_tensor(config))  # 构建vocab内的词嵌入矩阵\n",
    "\n",
    "    return train_loader, dev_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vWQU-wchE7Jz"
   },
   "outputs": [],
   "source": [
    "def extract_vocab_tensor(config):  # 提取vocab内的预训练词向量\n",
    "    \"\"\"\n",
    "    函数说明： 根据vocab内的字及索引构建词嵌入numpy(vocab是根据训练数据构建的词与索引对应的字典)\n",
    "\n",
    "    Parameter：\n",
    "    ----------\n",
    "        config.embedding_type - 'random' or others\n",
    "        config.vocab_path - vocab存储路径\n",
    "        config.pretrain_dir - 预训练词嵌入向量\n",
    "        config.embedding_dim - 词嵌入维度\n",
    "    Return:\n",
    "    -------\n",
    "        embedding_pretrained - 由vocab内的词组成的词嵌入数组（numpy）\n",
    "    \"\"\"\n",
    "    if config.embedding_type == 'random':  # 随机初始化\n",
    "        embedding_pretrained = None\n",
    "    else:  # 加载预训练词向量\n",
    "        vocab_tensor_path = config.embedding_type\n",
    "        if os.path.exists(vocab_tensor_path):  # 已构建则直接加载\n",
    "            embedding_pretrained = np.load(vocab_tensor_path)['embeddings'].astype('float32')\n",
    "        else:  # 重新构建\n",
    "            with open(config.vocab_path, 'rb') as vocab_f:\n",
    "                word_to_id = pickle.load(vocab_f)\n",
    "                pretrain_f = open(config.pretrain_dir, 'r', encoding='utf-8')\n",
    "                embeddings = np.random.rand(len(word_to_id), config.embedding_dim)\n",
    "                for i, line in enumerate(pretrain_f.readlines()):\n",
    "                    if i == 0:  # 若第一行是标题，则跳过 部分预训练模型第一行是词数和词嵌入维度\n",
    "                        continue\n",
    "                    lin = line.strip().split(' ')\n",
    "                    if lin[0] in word_to_id:\n",
    "                        idx = word_to_id[lin[0]]\n",
    "                        emb = [float(x) for x in lin[1:config.embedding_dim+1]]  # 取出对应词向量\n",
    "                        embeddings[idx] = np.asarray(emb, dtype='float')\n",
    "                pretrain_f.close()\n",
    "                np.savez_compressed(vocab_tensor_path, embeddings=embeddings)  # embeddings 是数组名\n",
    "                embedding_pretrained = embeddings.astype('float32')\n",
    "    return embedding_pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UR83GYkvEiPh",
    "outputId": "79c82905-5910-4a50-805d-5d7c4902b6ed"
   },
   "outputs": [],
   "source": [
    "config = Config()  # 实例化各参数\n",
    "train_loader, dev_loader, test_loader = build_dataset(config)  # 构造可用于模型输入的训练、验证、测试数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BT8qgmMBExDu",
    "outputId": "62c9ae82-008d-4b16-9467-cde663196fbe"
   },
   "outputs": [],
   "source": [
    "next(iter(train_loader))\n",
    "# length for \n",
    "# packed_inp = nn.utils.rnn.pack_padded_sequence(embed_with_char, lengths, batch_first=True)\n",
    "# packed_output, _ = self.lstm(packed_inp)\n",
    "# outputs, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "z1QWrKfCFpP4",
    "outputId": "c0ca510b-19d2-4c9f-eb42-cd8cb182236e"
   },
   "outputs": [],
   "source": [
    "!pip install pytorch-crf "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sCdb5eXeFXiy"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchcrf import CRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rawpfNHDF4fa"
   },
   "outputs": [],
   "source": [
    "START_TAG = 'START'\n",
    "STOP_TAG = 'STOP'\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(Model, self).__init__()\n",
    "        # if config.embedding_pretrained is not None:\n",
    "        #     self.embedding = nn.Embedding.from_pretrained(config.embedding_pretrained,\n",
    "        #                                                   freeze=False)  # 表示训练过程词嵌入向量会更新\n",
    "        # else:\n",
    "        self.embedding = nn.Embedding(config.vocab_len, config.embedding_dim)  # PAD索引填充\n",
    "        if config.bidirectional:\n",
    "            self.num_directions = 2\n",
    "        else:\n",
    "            self.num_directions = 1\n",
    "        self.config = config\n",
    "        self.rnn = nn.LSTM(config.embedding_dim, config.hidden_size, config.num_layers, batch_first=True,\n",
    "                           bidirectional=config.bidirectional)\n",
    "\n",
    "        self.tag_ls = self.getTagLs(config)\n",
    "        self.tag2idx = self.getTagDic()\n",
    "        # 转换参数矩阵 输入i,j是得分从j转换到i\n",
    "        self.tagset_size = len(self.tag2idx)\n",
    "        # 将lstm的输出映射到标记空间\n",
    "        self.hidden2tag = nn.Linear(config.hidden_size*self.num_directions, self.tagset_size)  \n",
    "        self.crf = CRF(self.tagset_size, batch_first=True)\n",
    "\n",
    "        self.tag_pad_idx = label_pad_id\n",
    "        \n",
    "        self.attn = nn.MultiheadAttention(\n",
    "            embed_dim=config.hidden_size * 2,\n",
    "            num_heads=4,\n",
    "            dropout=0.1\n",
    "        )\n",
    "        self.fc_dropout = nn.Dropout(0.1)\n",
    "\n",
    "\n",
    "    def _get_lstm_features(self, x):\n",
    "        # x:b_size sent\n",
    "        x = self.embedding(x)  # B -> (B, sent, e_d)\n",
    "        out, (hidden, c) = self.rnn(x)  # out:(B, sent, num_directions*hidden_size) hidden:(num_layer*nun_directions, B,  hidden_size)\n",
    "        out, attn_weight = self.attn(out, out, out)\n",
    "        \n",
    "        out = self.hidden2tag(self.fc_dropout(out))  # (B,num_directions*hidden_size) -> (B, num_class)\n",
    "        return out\n",
    "\n",
    "    def neg_log_likelihood(self, x, tags):  # 损失函数\n",
    "        feats = self._get_lstm_features(x)\n",
    "        return -self.crf(feats, tags)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x:b_size\n",
    "        lstm_feats = self._get_lstm_features(x)  # 获取BiLSTM的emission分数\n",
    "        out = self.crf.decode(lstm_feats)\n",
    "        return out\n",
    "\n",
    "    \n",
    "\n",
    "    def getTagLs(self, config):\n",
    "        tag_ls = config.class_ls\n",
    "        # tag_ls.append(START_TAG)\n",
    "        # tag_ls.append(STOP_TAG)\n",
    "        return tag_ls\n",
    "\n",
    "    def getTagDic(self):\n",
    "        tag_dic = {}\n",
    "        for idx, label in enumerate(self.tag_ls):\n",
    "            tag_dic[label] = idx\n",
    "        return tag_dic\n",
    "\n",
    "    def idx2Tag(self, idx):\n",
    "        return self.tag_ls[idx]\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Jl7wzogzmlDo"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Tp86HesJG-oN",
    "outputId": "06194d73-2e6d-4479-e389-4287591d31dc"
   },
   "outputs": [],
   "source": [
    "model = Model(config).to(config.device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TkFQhpMqHFSx"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import time\n",
    "from datetime import timedelta\n",
    "import numpy as np\n",
    "from sklearn import metrics\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mNWkODpeHMm_"
   },
   "outputs": [],
   "source": [
    "def accuracy( preds, y):\n",
    "    flatten_preds = [pred for sent_pred in preds for pred in sent_pred]\n",
    "    flatten_y = [tag for sent_tag in y for tag in sent_tag]\n",
    "    correct = [pred == tag for pred, tag in zip(flatten_preds, flatten_y)]\n",
    "    return sum(correct) / len(correct) if len(correct) > 0 else 0\n",
    "  \n",
    "\n",
    "\n",
    "def train_test(config, model, train_loader, dev_loader, test_loader):\n",
    "    start = time.time()\n",
    "\n",
    "    # 损失函数和优化器\n",
    "    # criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)\n",
    "    # 学习率指数衰减  每个epoch: lr = gamma*lr\n",
    "    scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)\n",
    "\n",
    "    create_dir_not_exists(config.model_save_dir)\n",
    "    if os.path.exists(config.model_save_name):  # 存在则加载模型 并继续训练\n",
    "        ckpt = torch.load(config.model_save_name)\n",
    "        model.load_state_dict(ckpt['model'])\n",
    "        optimizer.load_state_dict(ckpt['optimizer'])\n",
    "        start_epoch = ckpt['epoch']\n",
    "        max_acc = ckpt['dev_acc']\n",
    "        best_epoch = start_epoch\n",
    "        print(f'Load epoch {start_epoch} successful...')\n",
    "    else:\n",
    "        start_epoch = 0\n",
    "        max_acc = float('-inf')  # 负无穷\n",
    "        best_epoch = start_epoch\n",
    "    for epoch in range(start_epoch + 1, config.num_epochs + 1):\n",
    "        # train and evaluate\n",
    "        print('Epoch [{}/{}]'.format(epoch, config.num_epochs))\n",
    "        for i, data in enumerate(train_loader):\n",
    "            model.train()  # 开启dropout和BatchNormal\n",
    "            x, labels, lengths = data   \n",
    "            x, labels = x.to(config.device), labels.to(config.device)\n",
    "            predict = model(x)  # forward\n",
    "            # loss = criterion(predict, labels)\n",
    "            loss = model.neg_log_likelihood(x, labels)\n",
    "            optimizer.zero_grad()  # 梯度清零\n",
    "            loss.backward()  # backward\n",
    "            optimizer.step()  # update\n",
    "            if (i + 1) % 5 == 0:\n",
    "                # predict_labels = torch.max(predict, dim=1)[1].cpu().numpy()  # 如果在gpu上训练必须先转到CPU内才能转化为np类型\n",
    "                true = labels.data.cpu()\n",
    "                 \n",
    "                true_tags_list = [\n",
    "                                  [tag for tag in sent_tag  ] for sent_tag in true.tolist()\n",
    "                                  ]\n",
    "                train_acc = accuracy(predict, true_tags_list)\n",
    "                print(predict[0])\n",
    "                print(true_tags_list[0])\n",
    "                      \n",
    "\n",
    "                dev_acc, dev_loss = evaluate(config, dev_loader, model)\n",
    "                if dev_acc > max_acc:  # 存储最佳dev acc模型\n",
    "                    max_acc = dev_acc\n",
    "                    best_epoch = epoch  # 记录最佳轮数\n",
    "                    state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(),\n",
    "                             'epoch': epoch,'dev_acc': dev_acc}\n",
    "                    torch.save(state, config.model_save_name)\n",
    "                    improve = '*'  # 在有提升的结果后面加上*标注\n",
    "                else:\n",
    "                    improve = ''\n",
    "                msg = 'Iter: {0:>6}, Train loss: {1:>5.2},' \\\n",
    "                      ' Train acc: {2:>6.2%}, dev loss: {3:>5.2}, dev acc: {4:>6.2%}, Time: {5} {6}'\n",
    "                print(msg.format(i + 1, loss.item(), train_acc, dev_loss, dev_acc, time_since(start), improve))\n",
    "\n",
    "        # scheduler.step()  # lr衰减\n",
    "        if epoch - best_epoch >= config.require_improvement:\n",
    "            print(f'No optimization for {config.require_improvement} epoch, auto-stopping...')\n",
    "            break\n",
    "    print('Best model at epoch {}, dev acc: {:6.2%}, '.format(best_epoch, max_acc))\n",
    "    test(config, model, test_loader)  # 全程只进行一次测试\n",
    "    print('Time usage:', time_since(start))\n",
    "\n",
    "\n",
    "def evaluate(config, data_loader, model, test=False):\n",
    "    model.eval()  # 关闭dropout和BatchNormal\n",
    "    labels_all = np.array([], dtype=int)\n",
    "    predict_all = np.array([], dtype=int)\n",
    "    loss_total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in data_loader:\n",
    "            x, labels,lengths = data\n",
    "            x, labels = x.to(config.device), labels.to(config.device)\n",
    "\n",
    "            predict = model(x)  # forward\n",
    "            # loss = criterion(predict, labels)\n",
    "            loss = model.neg_log_likelihood(x, labels)\n",
    "            loss_total += loss.item()\n",
    "            labels = labels.cpu().numpy()\n",
    "            # predict_labels = torch.max(predict, dim=1)[1].cpu().numpy()\n",
    "            labels_all = np.append(labels_all, labels)  # 正确标签\n",
    "            predict_all = np.append(predict_all, predict)  # 预测标签\n",
    "    acc = metrics.accuracy_score(labels_all, predict_all)\n",
    "    return acc, loss_total / len(data_loader)  # 为何这里要除以len(data_loader) -> 因为这里处理了一个epoch的数据 而训练集输出loss只是处理一个batch_size的数据\n",
    "\n",
    "\n",
    "def test(config, model, test_loader):\n",
    "    ckpt = torch.load(config.model_save_name)\n",
    "    model.load_state_dict(ckpt['model'])\n",
    "    evaluate(config, test_loader, model, test=True)\n",
    "\n",
    "\n",
    "def create_dir_not_exists(filename):\n",
    "    if not os.path.exists(filename):\n",
    "        os.mkdir(filename)\n",
    "\n",
    "def time_since(since):\n",
    "    time_dif = time.time() - since\n",
    "    time_usage = timedelta(seconds=int(round(time_dif)))\n",
    "    return time_usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Zj77Y479HYiB",
    "outputId": "64b89e39-b2aa-4c80-a03e-2bcb43d750fb"
   },
   "outputs": [],
   "source": [
    "train_test(config, model, train_loader, dev_loader, test_loader)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
